import numpy as np

from htm.ApicalTMPairRegion import ApicalTMPairRegion
from htm.LocationGuassian2D import LocationGuassian2D
from htm.LocationSuperficial2D import LocationSupeficial2D


class L46CorticalColumn:

    def __init__(self, params, locationModuleParams, l4Params):
        # Paramters
        self.bumpType = params['bumpType']
        self.numOfLocModules = params['numOfLocModules']
        self.thresholds = params['thresholds']
        # Initialization
        if self.bumpType == "gaussian":
            self.L6aModules = [LocationGuassian2D(config) for config in locationModuleParams]
        elif self.bumpType == "square":
            self.L6aModules = [LocationSupeficial2D(config) for config in locationModuleParams]
        self.L4 = ApicalTMPairRegion(l4Params)
        self.l4CellCount = self.L4.columnCount * self.L4.cellsPerColumn

    def sensoryCompute(self, feedforward):
        inputParams = {
            "feedforward": feedforward,
            "basalInput": self.getLocationRepresentation(),
            "apicalInput": np.array([]),
            "basalGrowthCandidates": self.getLearnableLocationRepresentation(),
            "apicalGrowthCandidates": None
        }
        self.L4.compute(**inputParams)
        locationParams = {
            "anchorInput": self.L4.activeCells,
            "anchorGrowthCandidates": self.L4.winnerCells
        }
        for module in self.L6aModules:
            module.sensoryCompute(**locationParams)

    def activateRandomLocation(self):
        for module in self.L6aModules:
            module.activateRandomLocation()

    def reset(self):
        self.L4.reset()
        for module in self.L6aModules:
            module.reset()

    def setLearnMode(self, isLearning):
        self.L4.learn = isLearning
        for module in self.L6aModules:
            module.learn = isLearning

    def getLocationRepresentation(self):
        activeCells = np.array([], dtype="uint32")
        totalPrevCells = 0
        for module in self.L6aModules:
            activeCells = np.append(activeCells, module.activeCells + totalPrevCells)
            totalPrevCells += module.getNumOfCells()
        return activeCells

    def getLearnableLocationRepresentation(self):
        learnableCells = np.array([], dtype="uint32")
        totalPrevCells = 0
        for module in self.L6aModules:
            learnableCells = np.append(learnableCells, module.getLearnableCells() + totalPrevCells)
            totalPrevCells += module.getNumOfCells()
        return learnableCells

    def getSensoryAssociatedLocationRepresentation(self):
        cells = np.array([], dtype="uint32")
        totalPrevCells = 0
        for module in self.L6aModules:
            cells = np.append(cells, module.sensoryAssociatedCells + totalPrevCells)
            totalPrevCells += module.getNumOfCells()
        return cells

    def movementCompute(self, displacement, noiseFactor=0, moduleNoiseFactor=0):
        if noiseFactor != 0:
            xdisp = np.random.normal(0, noiseFactor)
            ydisp = np.random.normal(0, noiseFactor)
        else:
            xdisp = 0
            ydisp = 0
        locationParams = {
            "displacement": [displacement["top"] + ydisp, displacement["left"] + xdisp],
            "noiseFactor": moduleNoiseFactor
        }
        for module in self.L6aModules:
            module.movementCompute(**locationParams)

    def getLocationCopy(self):
        active_cells_list = []
        for module in self.L6aModules:
            active_cells_list.append(module.getActiveCells())
        return active_cells_list
